#include "cuda_img_resize.cuh"
//#define _DEBUG

#define CUDA_IMG_RESIZE_THREAD_NUM	1024
#define CUDA_IMG_RESIZE_WARP_SIZE	32

// image.h
typedef struct st_image {
	int w;
	int h;
	int c;
	float* data;
} image;

void CuImgResizeMalloc(st_cuda_resize_dev_space& space, int malloc_img_size)
{
	if (space.img_size < malloc_img_size)
	{
		CuImgResizeFree(space);

		space.img_size = malloc_img_size;

		cudaMalloc((void**)&space.img_rgb8, malloc_img_size * sizeof(unsigned char) * 3);
		cudaMalloc((void**)&space.img_rgb16, malloc_img_size * sizeof(unsigned short) * 3);
		cudaMalloc((void**)&space.img_rgbf, malloc_img_size * sizeof(float) * 3);
		cudaMalloc((void**)&space.img_rgbf_part, malloc_img_size * sizeof(float) * 3);
		cudaMalloc((void**)&space.img_rgbf_resized, malloc_img_size * sizeof(float) * 3);
	}
}

void CuImgResizeFree(st_cuda_resize_dev_space& space)
{
	cudaFree(space.img_rgb8);
	cudaFree(space.img_rgb16);
	cudaFree(space.img_rgbf);
	cudaFree(space.img_rgbf_part);
	cudaFree(space.img_rgbf_resized);
}

__device__ float get_pixel(image m, int x, int y, int c)
{
	//assert(x < m.w && y < m.h && c < m.c);
	return m.data[c * m.h * m.w + y * m.w + x];
}

__device__ void set_pixel(image m, int x, int y, int c, float val)
{
	if (x < 0 || y < 0 || c < 0 || x >= m.w || y >= m.h || c >= m.c) return;
	//assert(x < m.w && y < m.h && c < m.c);
	m.data[c * m.h * m.w + y * m.w + x] = val;
}
__device__ void add_pixel(image m, int x, int y, int c, float val)
{
	//assert(x < m.w && y < m.h && c < m.c);
	m.data[c * m.h * m.w + y * m.w + x] += val;
}

__global__ void KernelResizePart1(st_image im, st_image im_part, int w, int h)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	if (threadId >= im_part.w * im_part.h)
	{
		return;
	}

	int r, c, k;
	float w_scale = (float)(im.w - 1) / (w - 1);
	//float h_scale = (float)(im.h - 1) / (h - 1);

	r = threadId / im_part.w;
	c = threadId - r * im_part.w;

	for (k = 0; k < im.c; ++k)
	{
		float val = 0;
		if (c == w - 1 || im.w == 1)
		{
			val = get_pixel(im, im.w - 1, r, k);
		}
		else
		{
			float sx = c * w_scale;
			int ix = (int)sx;
			float dx = sx - ix;
			val = (1 - dx) * get_pixel(im, ix, r, k) + dx * get_pixel(im, ix + 1, r, k);
		}
		set_pixel(im_part, c, r, k, val);
	}
}

__global__ void KernelResizePart2(st_image im, st_image im_part, st_image im_resize, int w, int h)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	if (threadId >= im_resize.w * im_resize.h)
	{
		return;
	}

	int r, c, k;
	//float w_scale = (float)(im.w - 1) / (w - 1);
	float h_scale = (float)(im.h - 1) / (h - 1);

	r = threadId / im_resize.w;
	c = threadId - r * im_resize.w;

	//	int total_size = w*h;
	for (k = 0; k < im.c; ++k)
	{
		float sy = r * h_scale;
		int iy = (int)sy;
		float dy = sy - iy;

		float val = (1 - dy) * get_pixel(im_part, c, iy, k);
		set_pixel(im_resize, c, r, k, val);

		if (r == h - 1 || im.h == 1)
		{
			val = get_pixel(im_part, c, iy, k);
			set_pixel(im_resize, c, h - 1, k, val);
		}
		else
		{
			val = dy * get_pixel(im_part, c, iy + 1, k);
			add_pixel(im_resize, c, r, k, val);
		}
	}
}

__global__ void KernelInt8ToFloat(unsigned char* img_rgb8, float* img_t, int w, int h)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	unsigned int img_size = w * h;
	if (threadId >= img_size)
	{
		return;
	}
	img_t[threadId] = float(img_rgb8[threadId]) / 255.f;
	img_t[threadId + img_size] = float(img_rgb8[threadId + img_size]) / 255.f;
	img_t[threadId + img_size * 2] = float(img_rgb8[threadId + img_size * 2]) / 255.f;
}

__global__ void KernelInt16ToFloat(unsigned short* img_rgb16, float* img_t, int w, int h)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	unsigned int img_size = w * h;
	if (threadId >= img_size)
	{
		return;
	}
	img_t[threadId] = float(img_rgb16[threadId]) / 65535.f;
	img_t[threadId + img_size] = float(img_rgb16[threadId + img_size]) / 65535.f;
	img_t[threadId + img_size * 2] = float(img_rgb16[threadId + img_size * 2]) / 65535.f;
}

__global__ void KernelMean(float* img_t, int w, int h, float mean_r, float mean_g, float mean_b)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	unsigned int img_size = w * h;
	if (threadId >= img_size)
	{
		return;
	}
	img_t[threadId] = img_t[threadId] - mean_r;
	img_t[threadId + img_size] = img_t[threadId + img_size] - mean_g;
	img_t[threadId + img_size * 2] = img_t[threadId + img_size * 2] - mean_b;
}

__global__ void KernelDev(float* img_t, int w, int h, float dev_r, float dev_g, float dev_b)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	unsigned int img_size = w * h;
	if (threadId >= img_size)
	{
		return;
	}
	img_t[threadId] = img_t[threadId] / dev_r;
	img_t[threadId + img_size] = img_t[threadId + img_size] / dev_g;
	img_t[threadId + img_size * 2] = img_t[threadId + img_size * 2] / dev_b;
}

void CuImgResizeOp(int w, int h, st_cuda_resize_dev_space& space, int w_r, int h_r)
{
	st_image im_orig, im_part, im_resize;

	im_orig.w = w;
	im_orig.h = h;
	im_orig.c = 3;
	im_orig.data = space.img_rgbf;

	im_part.w = w_r;
	im_part.h = h;
	im_part.c = 3;
	im_part.data = space.img_rgbf_part;

	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((w_r * h + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelResizePart1 << <blocks, threads >> > (im_orig, im_part, w_r, h_r);
	}
	//cudaDeviceSynchronize();

	im_resize.w = w_r;
	im_resize.h = h_r;
	im_resize.c = 3;
	im_resize.data = space.img_rgbf_resized;

	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((w_r * h_r + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelResizePart2 << <blocks, threads >> > (im_orig, im_part, im_resize, w_r, h_r);
	}
	//cudaDeviceSynchronize();
}

__global__ void KernelCopyMakeBorder(float* im_src, int src_w, int src_h, float* im_dest, int dest_w, int dest_h, int left, int top)
{
	int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	int dest_img_size = dest_w * dest_h;
	if (threadId >= dest_img_size)
	{
		return;
	}
	int x = threadId % dest_w - left;
	int y = threadId / dest_w - top;
	int src_img_size = src_w * src_h;

	if (x >= 0 && x < src_w && y >= 0 && y < src_h)
	{
		int index = x + y * src_w;

		im_dest[threadId] = im_src[index];

		threadId += dest_img_size;
		index += src_img_size;
		im_dest[threadId] = im_src[index];

		threadId += dest_img_size;
		index += src_img_size;
		im_dest[threadId] = im_src[index];
	}
	else
	{
		im_dest[threadId] = 0.f;
		threadId += dest_img_size;
		im_dest[threadId] = 0.f;
		threadId += dest_img_size;
		im_dest[threadId] = 0.f;
	}
}

void CuImgCopyMakeBorder(int resize_w, int resize_h, st_cuda_resize_dev_space& space, int net_w, int net_h)
{
	st_image im_resize, im_net;
	im_resize.w = resize_w;
	im_resize.h = resize_h;
	im_resize.c = 3;
	im_resize.data = space.img_rgbf_resized;

	im_net.w = net_w;
	im_net.h = net_h;
	im_net.c = 3;
	im_net.data = space.img_rgbf_net;


	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((net_w * net_h + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelCopyMakeBorder << <blocks, threads >> > (
			im_resize.data, im_resize.w, im_resize.h,
			im_net.data, im_net.w, im_net.h,
			0, 0);
	}
}

void CuImgMeanDev(st_cuda_resize_dev_space& space, int w_r, int h_r, float* mean, float* dev)
{
	dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
	dim3 blocks = dim3((w_r * h_r + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
	KernelMean << <blocks, threads >> > (space.img_rgbf_resized, w_r, h_r, mean[0], mean[1], mean[2]);
	KernelDev << <blocks, threads >> > (space.img_rgbf_resized, w_r, h_r, dev[0], dev[1], dev[2]);
}


void CuImgPreProcess(int w, int h,
	int resize_w, int resize_h,
	st_cuda_resize_dev_space& space, int net_w, int net_h,
	float* mean, float* dev)
{
	CuImgResizeOp(w, h, space, resize_w, resize_h);
	CuImgMeanDev(space, resize_w, resize_h, mean, dev);
	CuImgCopyMakeBorder(resize_w, resize_h, space, net_w, net_h);
}

void CuImgResizeRgb8(unsigned char** host_img_rgb8, int w, int h,
	int resize_w, int resize_h,
	st_cuda_resize_dev_space& space, int net_w, int net_h,
	float* mean, float* dev)
{
	cudaError_t error; //store cuda error codes
	for (int i = 0; i < 3; ++i)
	{
		error = cudaMemcpy(space.img_rgb8 + i * w * h, host_img_rgb8[i], w * h * sizeof(unsigned char), cudaMemcpyHostToDevice);

#ifdef _DEBUG
		if (error != cudaSuccess)
		{
			printf("cudaMemcpy, returned error %s (code %d), line(%d)\n", cudaGetErrorString(error), error, __LINE__);
			exit(EXIT_FAILURE);
		}
#endif
	}
	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((w * h + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelInt8ToFloat << <blocks, threads >> > (space.img_rgb8, space.img_rgbf, w, h);
		//cudaDeviceSynchronize();
	}
	CuImgPreProcess(w, h, resize_w, resize_h, space, net_w, net_h, mean, dev);
}


void CuImgResizeRgb16(unsigned short** host_img_rgb16, int w, int h,
	int resize_w, int resize_h,
	st_cuda_resize_dev_space& space, int net_w, int net_h,
	float* mean, float* dev)
{
	cudaError_t error; //store cuda error codes
	for (int i = 0; i < 3; ++i)
	{
		error = cudaMemcpy(space.img_rgb16 + i * w * h, host_img_rgb16[i], w * h * sizeof(unsigned short), cudaMemcpyHostToDevice);

#ifdef _DEBUG
		if (error != cudaSuccess)
		{
			printf("cudaMemcpy, returned error %s (code %d), line(%d)\n", cudaGetErrorString(error), error, __LINE__);
			exit(EXIT_FAILURE);
		}
#endif
	}
	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((w * h + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelInt16ToFloat << <blocks, threads >> > (space.img_rgb16, space.img_rgbf, w, h);
	}

	CuImgPreProcess(w, h, resize_w, resize_h, space, net_w, net_h, mean, dev);
}

void CuImgResizeRgbF(float** host_img_rgbf, int w, int h,
	int resize_w, int resize_h,
	st_cuda_resize_dev_space& space, int net_w, int net_h,
	float* mean, float* dev)
{
	cudaError_t error; //store cuda error codes

	for (int i = 0; i < 3; ++i)
	{
		error = cudaMemcpy(space.img_rgbf + i * w * h, host_img_rgbf[i], w * h * sizeof(float), cudaMemcpyHostToDevice);

#ifdef _DEBUG
		if (error != cudaSuccess)
		{
			printf("cudaMemcpy, returned error %s (code %d), line(%d)\n", cudaGetErrorString(error), error, __LINE__);
			exit(EXIT_FAILURE);
		}
#endif
	}
	CuImgPreProcess(w, h, resize_w, resize_h, space, net_w, net_h, mean, dev);
}